#include <adjust_result.h>
// 计算重叠面积比例
double calculateOverlap(const std::vector<double> &box1, const std::vector<double> &box2) {
    double xOverlap = std::max(0.0, std::min(box1[0] + box1[2] / 2, box2[0] + box2[2] / 2) - std::max(box1[0] - box1[2] / 2, box2[0] - box2[2] / 2));
    double yOverlap = std::max(0.0, std::min(box1[1] + box1[3] / 2, box2[1] + box2[3] / 2) - std::max(box1[1] - box1[3] / 2, box2[1] - box2[3] / 2));
    return 1 - xOverlap * yOverlap / (box1[2] * box1[3] + box2[2] * box2[3] - xOverlap * yOverlap);
}

// 识别框的6要素
struct BoundingBox {
    double x_center;
    double y_center;
    double width;
    double height;
    double radius;
    double colonum;
};
// 提取框的6要素
std::vector<BoundingBox> parseBoundingBoxes(const std::vector<std::vector<double>> &data) {
    std::vector<BoundingBox> boxes;
    
    for (const std::vector<double> &boxData : data) {
        BoundingBox box;
        box.x_center = boxData[0];
        box.y_center = boxData[1];
        box.width = boxData[2];
        box.height = boxData[3];
        box.radius = boxData[4];
        box.colonum = boxData[5];
        
        boxes.push_back(box);
    }
    
    return boxes;
}

// 比较函数，用于指定排序规则，按照第5列的值进行排序
bool compareRows(const std::vector<double>& row1, const std::vector<double>& row2) {
    return row1[5] < row2[5];  // 从小到大排序
}

// 转置函数
std::vector<std::vector<float>> transpose(const std::vector<std::vector<float>>& matrix) {
    int rows = matrix.size();
    int cols = matrix[0].size();

    std::vector<std::vector<float>> transposed(cols, std::vector<float>(rows));

    for (int y = 0; y < rows; ++y) {
        for (int x = 0; x < cols; ++x) {
            transposed[x][y] = matrix[y][x];
        }
    }

    return transposed;
}

// 筛选缺陷
std::vector<std::vector<double>> select(const std::vector<std::vector<float>>& data, float threshold, int rows, int cols) {
    std::vector<std::vector<double>> filtered_data;

    for (int i = 0; i < rows; ++i) {
        bool exceedsThreshold = false;
        for (int j = 4; j < cols; ++j)
            if (data[i][j] > threshold) {
                exceedsThreshold = true;
                break;
            }

        if (exceedsThreshold) {
            std::vector<double> filtered_row(6, 0.0);
            for (int j = 0; j < 4; ++j) 
                filtered_row[j] = static_cast<int>(data[i][j]);  // 取整存储到新数组中

            for (int j = 4; j < cols; ++j) 
                if (data[i][j] > threshold) {
                    filtered_row[4] = data[i][j];  // 存储大于阈值的值
                    filtered_row[5] = j-4;  // 大于阈值的列数
                    break;
                }

            filtered_data.push_back(filtered_row);
        }
    }

    // 使用比较函数来对二维数组按照第5列数据进行排序
    std::sort(filtered_data.begin(), filtered_data.end(), compareRows);

    return filtered_data;
}

// 恢复原图缺陷坐标
std::vector<std::vector<double>> return_(const std::vector<std::vector<double>>& data, int h, int w){
    std::vector<std::vector<double>> filtered_data = data;
    if (h > w) {
        double scale = static_cast<double>(h) / 640.0;
        for (int i = 0; i < filtered_data.size(); ++i) {
            for (int j = 0; j < 4; ++j) 
                filtered_data[i][j] *= scale;  // 缩放前4列
            filtered_data[i][0] -= (h / 2 - w / 2);  // 修改第0列
        }
    } else{
        double scale = static_cast<double>(w) / 640.0;
        for (int i = 0; i < filtered_data.size(); ++i) {
            for (int j = 0; j < 4; ++j) 
                filtered_data[i][j] *= scale;  // 缩放前4列
            filtered_data[i][1] -= (w / 2 - h / 2);  // 修改第1列
        }
    }

    return filtered_data;
}

// 删除重叠框
std::vector<std::vector<double>> nms_(const std::vector<std::vector<double>>& data, float threshold){
    std::vector<std::vector<double>> result = data;
    int i = 0;
    while (i < result.size() - 1) {
        if (result[i][5] != result[i + 1][5]){
            i = i + 1;
            continue;
        }
            
        double overlap = calculateOverlap(result[i], result[i + 1]);
        if (overlap < threshold && result[i][4] >= result[i + 1][4]) 
            result.erase(result.begin() + i + 1);
        else if (overlap < threshold && result[i][4] < result[i + 1][4])
            result.erase(result.begin() + i);
        else if (overlap >= threshold)
            i = i + 1;
    }

    return result;
}
// 绘制缺陷框
cv::Mat draw_image(const std::vector<std::vector<double>>& data, const cv::Mat &Image, const std::vector<std::string>& typeNames, const std::vector<std::vector<int>>& colors){
    cv::Mat inputImage = Image.clone();
    std::vector<BoundingBox> boxes = parseBoundingBoxes(data);
    for (const BoundingBox &box : boxes) {
        // 绘制方框
        int colorIndex = static_cast<int>(box.colonum);;
        cv::Scalar rectColor(colors[colorIndex][0], colors[colorIndex][1], colors[colorIndex][2]);
        cv::Point topLeft(box.x_center - box.width / 2, box.y_center - box.height / 2);
        cv::Point bottomRight(box.x_center + box.width / 2, box.y_center + box.height / 2);
        cv::rectangle(inputImage, topLeft, bottomRight, rectColor, 2);

        // 绘制缺陷种类和缺陷置信系数
        std::ostringstream doubleStream;
        doubleStream << box.radius;
        std::string doubleString = doubleStream.str();
        std::string combinedData = typeNames[colorIndex] + " " + doubleString;
        cv::Point textPosition(box.x_center - box.width / 2, box.y_center - box.height / 2);
        cv::putText(inputImage, combinedData, textPosition, cv::FONT_HERSHEY_SIMPLEX, 2.0, rectColor, 2);
    }

    return inputImage;
}